import time
from pymilvus import connections, Collection
from sentence_transformers import SentenceTransformer

# 连接 Milvus（使用与数据导入一致的连接名）
from config import Config
connections.connect(Config.MILVUS_CONN, host=Config.MILVUS_HOST, port=Config.MILVUS_PORT)

# 载入笑话集合（已在 Milvus 中建立）
try:
    col = Collection(Config.MILVUS_COLLECTION, using=Config.MILVUS_CONN)
    col.load()
except Exception:
    col = None

# 向量模型（常驻内存避免反复加载）
embedder = Config.get_embedder()  # 维度需与Milvus schema一致

def search_jokes(query: str, k: int = 4) -> list[str]:
    """从 Milvus 搜索相似笑话，返回至多k条文本列表"""
    print(f"Searching for jokes: {query} (k={k})")
    if not col:
        return []
    emb = embedder.encode([query])

    # 简单超时保护：Milvus SDK 没有原生超时，这里用轮询退避
    start = time.time()
    res = None
    while True:
        try:
            res = col.search(emb, "embedding", param={"nprobe": Config.MILVUS_NPROBE}, limit=max(k, 5), output_fields=["text"])
            break
        except Exception:
            if time.time() - start > 2.5:  # ~2.5s超时
                return []
            time.sleep(0.05)

    hits = res[0] if res else []
    jokes: list[str] = []
    for h in hits:
        txt = None
        try:
            if hasattr(h, "entity") and h.entity is not None:
                txt = h.entity.get("text")
        except Exception:
            txt = None
        if txt is None:
            txt = getattr(h, "text", None) or (h.get("text") if hasattr(h, "get") else None)
        if isinstance(txt, str) and txt.strip():
            jokes.append(txt.strip())
        if len(jokes) >= k:
            break
    return jokes

# 兼容旧接口：返回Top-1
def search_joke(query: str) -> str | None:
    jokes = search_jokes(query, k=1)
    return jokes[0] if jokes else None